#ifndef GEO_MF_SIMPLE_CC
#define GEO_MF_SIMPLE_CC
#include<string>
#include<fstream>
#include<sstream>
#include<iostream>
#include<algorithm>
#include<iterator>
#include<vector>
#include<map>
#include<cstdlib>
#include<ctime>
#include<cmath>
#include "def.hpp"
#include "utils.hpp"
#include "topicmodel_geo_abstract.hpp"

class GeoMFSimple:public GeoTopicModel{
protected:
  double* phi_0; // background language model for the whole corpus
  GeoTerms phi_geo;
  vector<int> doc_n;
  int* geo_docs;
  double* phi_norm;  

  double* phi_norm_buffer; // used for calculating phi norm
  double lowerbound;
  double learning_rate;
  double reg;
  int N;
  ofstream file_output;
  int K; // parameter for L1 regularization by using Truncated Gradient method
  double g; // parameter for controling the sparsity
public:
  GeoMFSimple();
  virtual void Inference();
  virtual void Prediction();
  virtual double Likelihood();
  virtual int ReadDOCS();
  virtual int GetDocuments(int); // this indicate how documents are read through into the model
  virtual int SaveModel();
private:
  int Single_M_Step_L2();
  int Single_M_Step_L1();
};

int GeoMFSimple::SaveModel(){
  ofstream outFile;
  string  model_file = PuzaDEF::Instance()->output_file_name + ".phi0";
  outFile.open(model_file.c_str());
  for (int i=0; i<PuzaDEF::Instance()->TERM_NUM; i++){
    outFile << i << ":" << phi_0[i] << " " << endl;;
  }
  outFile.flush();
  outFile.close();

  model_file = PuzaDEF::Instance()->output_file_name + ".phi_geo";
  outFile.open(model_file.c_str());
  for(int r=0; r<PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    outFile << geo_mapping[r] << "\t";
    for(map<int,double>::iterator iter = phi_geo[r].begin(); iter != phi_geo[r].end(); iter++){
      outFile << (*iter).first << ":" << (*iter).second << " ";
    }
    outFile << endl;
  }
  outFile.flush();
  outFile.close();
  return 0;
}

GeoMFSimple::GeoMFSimple(){
  phi_0 = new double[PuzaDEF::Instance()->TERM_NUM];
  phi_norm = new double[PuzaDEF::Instance()->GEO_REGION_NUM];
  phi_norm_buffer = new double[PuzaDEF::Instance()->TERM_NUM];
  geo_docs = new int[PuzaDEF::Instance()->GEO_REGION_NUM];

  // some initialization
  for(int i=0; i < PuzaDEF::Instance()->TERM_NUM; i++){
    phi_0[i] = 0.0;
  }

  for(int r=0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    geo_docs[r] = 0;
  }

  phi_geo.resize(PuzaDEF::Instance()->GEO_REGION_NUM);

  learning_rate = PuzaDEF::Instance()->GEO_LEARNING_RATE;
  reg = PuzaDEF::Instance()->GEO_REG;
  K = PuzaDEF::Instance()->GEO_REG_L1_K;
  g = PuzaDEF::Instance()->GEO_REG_L1_G;
}

int GeoMFSimple::Single_M_Step_L2(){
  int flag = 0;

  for(unsigned int d = 0; d < TM_DOCS.size(); d++){
    int r = doc_geo_stamp[d];
    TermList current_list = TM_DOCS[d];
    int term_count = current_list.size();
      
    for (int j=0;j < term_count; j++){
      int term_id = current_list[j].term_id;
      int count = current_list[j].term_stat;
      double gradient_pi_k_w = count -  doc_n[d] * exp(phi_0[term_id] + phi_geo[r][term_id] - phi_norm[r]) - (reg / geo_docs[r]) * phi_geo[r][term_id];
      double phi_minus = phi_0[term_id] + phi_geo[r][term_id];
      phi_geo[r][term_id] = phi_geo[r][term_id] + learning_rate * gradient_pi_k_w;
      double phi_addition = phi_0[term_id] + phi_geo[r][term_id];
      phi_norm[r] = log_sub(phi_norm[r],phi_minus); // important! no minus sign
      phi_norm[r] = log_sum(phi_norm[r],phi_addition);
    }
  }
  return flag;
}

int GeoMFSimple::Single_M_Step_L1(){
  int flag = 0;

  for(unsigned int d = 0; d < TM_DOCS.size(); d++){
    int r = doc_geo_stamp[d];
    TermList current_list = TM_DOCS[d];
    int term_count = current_list.size();

    if( d % K != 0){
      // normal gradient descent
      for (int j=0;j < term_count; j++){
	int term_id = current_list[j].term_id;
	int count = current_list[j].term_stat;
	double gradient_pi_k_w = count -  doc_n[d] * exp(phi_0[term_id] + phi_geo[r][term_id] - phi_norm[r]);
	double phi_minus = phi_0[term_id] + phi_geo[r][term_id];
	phi_geo[r][term_id] = phi_geo[r][term_id] + learning_rate * gradient_pi_k_w;
	double phi_addition = phi_0[term_id] + phi_geo[r][term_id];
	phi_norm[r] = log_sub(phi_norm[r],phi_minus); // important! no minus sign
	phi_norm[r] = log_sum(phi_norm[r],phi_addition);
      }
    }
    else{
      // truncated gradient descent
      for (int j=0;j < term_count; j++){
	int term_id = current_list[j].term_id;
	int count = current_list[j].term_stat;
	double gradient_pi_k_w = count -  doc_n[d] * exp(phi_0[term_id] + phi_geo[r][term_id] - phi_norm[r]);
	double phi_minus = phi_0[term_id] + phi_geo[r][term_id];
	phi_geo[r][term_id] = Trunc(phi_geo[r][term_id] + learning_rate * gradient_pi_k_w, learning_rate * g * K); // truncated gradient step
	double phi_addition = phi_0[term_id] + phi_geo[r][term_id];
	phi_norm[r] = log_sub(phi_norm[r],phi_minus); // important! no minus sign
	phi_norm[r] = log_sum(phi_norm[r],phi_addition);
      }
    }
  }
  return flag;
}

void GeoMFSimple::Inference(){
  PuzaLogger::Instance()->PutString("Start Inference for Simple Geo Matrix Factorization Model");
  int ITER = 0;
  double old_lowerbound = 0;
  while(1){
    PuzaLogger::Instance()->PutString("Iteration "+to_string<int>(ITER,std::dec));
    if(PuzaDEF::Instance()->GEO_REG_TYPE == 0){
      Single_M_Step_L2();
    }
    else{
      Single_M_Step_L1();
    }
    // phi norm check
    /*    for(int k=0; k < PuzaDEF::Instance()->GEO_REGION_NUM; k++){
      for(int i=0; i < PuzaDEF::Instance()->TERM_NUM; i++){
	phi_norm_buffer[i] = 0.0;
	map<int,double>::iterator geo_term_iter =phi_geo[k].find(i);
	if(geo_term_iter != phi_geo[k].end()){
 phi_geo[r][term_id] = phi_geo[r][term_id] + learning_rate * gradient_pi_k_w;	   phi_norm_buffer[i] = phi_geo[k][i];
	} 
	phi_norm_buffer[i] = phi_norm_buffer[i] + phi_0[i];
      }
      double result = log_sum_array(phi_norm_buffer, PuzaDEF::Instance()->TERM_NUM);
      cout << result << " " << phi_norm[k] << " " << phi_norm[k] - result << endl;
    }
    exit(0);
    */

    Likelihood();
    PuzaLogger::Instance()->PutString("\tLikelihood:" + to_string<double>(lowerbound,std::dec));
    if((lowerbound < old_lowerbound) && ( ITER != 0))
      break;
    old_lowerbound = lowerbound;
    if (ITER>=500){
      break;
    }
    if (ITER % 50 == 0){
      SaveModel();
    }
    ITER ++;
  }
  PuzaLogger::Instance()->PutString("Batch Learning Finished.");
}

void GeoMFSimple::Prediction(){
}
 
double GeoMFSimple::Likelihood(){
  double result = 0.0;

  for(int r = 0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    for(int t=0; t < PuzaDEF::Instance()->TERM_NUM; t++){
      phi_norm_buffer[t] = phi_0[t] + phi_geo[r][t];
    }
    double temp = log_sum_array(phi_norm_buffer,PuzaDEF::Instance()->TERM_NUM);
    cout << temp << " " << phi_norm[r] << " " << phi_norm[r] - temp << endl;
  }
  exit(0);


  for(unsigned int d = 0; d < TM_DOCS.size(); d++){
    int current_region = doc_geo_stamp[d];
    TermList current_list = TM_DOCS[d];
    int term_count = current_list.size();
    for (int j=0;j < term_count; j++){
      int term_id = current_list[j].term_id;
      int count = current_list[j].term_stat;
      result = result + count * (phi_0[term_id] + phi_geo[current_region][term_id] - phi_norm[current_region]);
    }      
  }
  lowerbound = result;
  return lowerbound;
}

int GeoMFSimple::ReadDOCS(){
  int MAX_INT = 0;
  N = 0;
  string temp;
  if(PuzaDEF::Instance()->input_file_name.empty()){
    cout << "Please give input file name." << endl;
    exit(0);
  }
  PuzaDEF::Instance()->InputFile.open(PuzaDEF::Instance()->input_file_name.c_str());
  if(!PuzaDEF::Instance()->InputFile){
    cout << "Can't open the file." << endl;
    exit(0);
  }
  PuzaDEF::Instance()->data_ids.clear();
  string original_id;
  string temp_geo_stamp;
  string temp_user_stamp;
  vector<string> buffer;
  doc_n.clear();
  while(!PuzaDEF::Instance()->InputFile.eof()){
    getline(PuzaDEF::Instance()->InputFile,temp,'\n');
    std::istringstream iss(temp);
    // get the doc id
    getline(iss,temp,'\t');
    original_id = temp;
    // get geo                                                                                                                                 
    getline(iss,temp,'\t');
    temp_geo_stamp = temp;
    // get user
    getline(iss,temp,'\t');
    temp_user_stamp = temp;
    // get the doc body
    getline(iss,temp,'\n');
    std::istringstream tempiss(temp);
    vector<string> tokens;
    copy(istream_iterator<string>(tempiss), istream_iterator<string>(), back_inserter<vector<string> >(tokens));
    if(tokens.size()<1){
      continue;
    }

    int now_id = PuzaDEF::Instance()->data_ids.size();
    PuzaDEF::Instance()->data_ids[now_id] = original_id;
    map<string,int>::iterator geo_iter = geo_reverse_mapping.find(temp_geo_stamp);
    int geo_id = -1;
    if(geo_iter != geo_reverse_mapping.end()){
      // id found
      geo_id = (*geo_iter).second;
      doc_geo_stamp.push_back(geo_id);
    }
    else{
      geo_id = geo_reverse_mapping.size();
      geo_reverse_mapping[temp_geo_stamp] = geo_id;
      geo_mapping[geo_id] = temp_geo_stamp;
      doc_geo_stamp.push_back(geo_id);
    }
    TermList assignList;
    int temp_n = 0;
    for (vector<string>::iterator iter = tokens.begin(); iter < tokens.end(); iter++){
      int word_id;
      int word_count;
      std::istringstream valueiss((*iter));
      getline(valueiss,temp,':');
      from_string<int>(word_id,temp,std::dec);
      getline(valueiss,temp);  // get count
      from_string<int>(word_count,temp,std::dec);
      TermItem newAssign;
      newAssign.term_id = word_id;
      newAssign.term_stat = word_count;
      assignList.push_back(newAssign);
      phi_0[word_id] = phi_0[word_id] + word_count;
      map<int,double>::iterator geo_term_iter = phi_geo[geo_id].find(word_id);
      if(geo_term_iter == phi_geo[geo_id].end()){
	phi_geo[geo_id][word_id] = 0.0;
      }
      N = N + word_count;
      temp_n = temp_n + word_count;
    }
    if((int)assignList.size() > MAX_INT)
      MAX_INT = assignList.size();
    TM_DOCS.push_back(assignList);
    doc_n.push_back(temp_n);
    geo_docs[geo_id] = geo_docs[geo_id] + 1;
  }
  PuzaDEF::Instance()->InputFile.close();
  return MAX_INT;
}

int GeoMFSimple::GetDocuments(int choice){
  ReadDOCS();
  for(int i=0; i < PuzaDEF::Instance()->TERM_NUM; i++){
    phi_0[i] = log((phi_0[i] + 1.0)/(N+PuzaDEF::Instance()->TERM_NUM));
    phi_norm_buffer[i] = phi_0[i];
  }
  double phi_norm_temp = log_sum_array(phi_norm_buffer,PuzaDEF::Instance()->TERM_NUM);
  for(int k=0; k < PuzaDEF::Instance()->GEO_REGION_NUM; k++){
    phi_norm[k] = phi_norm_temp;
  }
  return 0;
}
#endif
